import pyb
import sensor, image, time, tf

# 初始化相机
sensor.reset()
sensor.set_pixformat(sensor.RGB565)  # 设置摄像头像素格式
sensor.set_framesize(sensor.QVGA)   # 设置摄像头分辨率
sensor.set_brightness(700)          # 设置摄像头亮度
sensor.skip_frames(time = 200)
clock = time.clock()

# 模型和标签路径
net_path = "/sd/moblie1.tflite"      # 模型文件路径
labels_path = "/sd/labels1.txt"      # 标签文件路径
labels = [line.rstrip() for line in open(labels_path)]  # 加载标签文件

# 大类和小类的映射
category_mapping = {
    "电子外设": ["mouse", "keyboard", "monitor", "headphones", "speaker", "printer", "mobile_phone"],
    "常用工具": ["wrench", "soldering_iron", "electrodrill", "tape_measure", "screwdriver", "pliers", "oscillograph", "multimeter"]
}

# 加载模型
try:
    net = tf.load(net_path, load_to_fb=True)  # 加载模型到帧缓冲区
    print("Model loaded successfully.")
except Exception as e:
    print("Failed to load model:", e)
    while True:
        pass

# 主循环
while True:
    clock.tick()  # 开始计时
    img = sensor.snapshot()  # 捕获图像帧

    # 增加对图像进行阈值处理来提高物体对比度，可能有助于提高目标检测的效果
    img.binary([([0, 100], [0, 100], [0, 100])])  # 可以调整阈值根据环境

    # 搜索图像中的矩形区域
    for r in img.find_rects(threshold=10000):  # 调整阈值以控制矩形检测灵敏度
        img.draw_rectangle(r.rect(), color=(255, 0, 0))  # 在图像中绘制矩形
        roi = r.rect()  # 获取矩形区域坐标 (x, y, w, h)
        img1 = img.copy(roi=roi)  # 拷贝矩形区域图像用于分类

        # 使用模型分类矩形区域
        for obj in tf.classify(net, img1, min_scale=1.0, scale_mul=0.5, x_overlap=0.25, y_overlap=0.25):
            x, y, w, h = obj.rect()  # 获取检测窗口位置
            img.draw_rectangle((x, y, w, h), color=(0, 255, 0))  # 绘制检测窗口
            sorted_list = sorted(zip(labels, obj.output()), key=lambda x: x[1], reverse=True)

            # 获取置信度最高的标签
            label, confidence = sorted_list[0]
            print("Label: {}, Confidence: {:.2f}".format(label, confidence))

            # 输出大类和小类信息
            big_category = None
            for category, items in category_mapping.items():
                if label in items:
                    big_category = category
                    break

            # 打印大类和小类
            if big_category:
                print(f"大类: {big_category}, 小类: {label}")
            else:
                print("未找到对应的大类和小类")

            # 在图像中显示结果
            img.draw_string(x, y - 10, "{}:{:.2f}".format(label, confidence), color=(0, 255, 0))

    # 打印帧率
    print("FPS: {:.2f}".format(clock.fps()))
